#pragma once

#include "geometry_msgs/msg/pose_array.hpp"
#include <opencv2/opencv.hpp>

#include "final_exam/armor.h"
#include "final_exam/kalman_predictor.h"
#include "final_exam/param.h"
#include "final_exam/pose_solver.h"

// 装甲板识别
class ArmorDetect {
private:
  bool enemy;                                          // 敌人颜色
  cv::Mat img;                                         // 输出图像
  cv::Mat src;                                         // 原始图像
  PoseSolver pose_solver;                              // 位置解算器
  std::vector<std::pair<Predictor, Armor>> predictors; // 位置预测器
  geometry_msgs::msg::PoseArray armor_poses_msg;       // 装甲板位置消息

  // BGR 转灰度图
  void BGR2Gray(cv::Mat &src, cv::Mat &gray);
  // 图像增强和二值化
  void imageEnhance(cv::Mat &gray, cv::Mat &enhance);
  // 轮廓检测，筛选出灯条
  void findLightBar(cv::Mat &enhance, std::vector<cv::RotatedRect> &light_bars);
  // 灯条配对，筛选出装甲板
  void matchLightBar(std::vector<cv::RotatedRect> &light_bars,
                     std::vector<Armor> &armors);
  // 画出装甲板，同时将位置信息存入 armor_poses_msg
  void drawArmor();
  // 辅助函数：根据中心点和象限重新排序角点
  void reorderArmorPoints(const std::vector<cv::Point2f> &input_points,
                          std::vector<cv::Point2f> &output_points);
  // 解算装甲板位置
  void solveArmorPose(std::vector<Armor> &armors);
  // 获取装甲板图像
  void getArmorImg(const cv::Mat &src, const Armor &armor, cv::Mat &armor_img);
  // 预测下一位置
  void kalmanPredict();
  // 卡尔曼滤波修正当前位置
  void kalmanCorrect(cv::Point3f &pt, const Armor &armor);
  // 将 tvec/rvec 转换为 geometry_msgs::msg::Pose（位置 + 四元数）
  geometry_msgs::msg::Pose poseFromRt(const cv::Mat &tvec, const cv::Mat &rvec);

public:
  ArmorDetect(bool enemy = param::RED) : enemy(enemy) {}
  // 设置敌人颜色
  void setEnemy(bool enemy);
  // 灯条检测主函数
  void detect(cv::Mat &src, cv::Mat &dst);
  // 设置相机参数
  void setCameraParams(const cv::Mat &camera_matrix,
                       const cv::Mat &dist_coeffs);
  // 设置物体点
  void setObjectPoints(const std::vector<cv::Point3f> &object_points);
  // 加载兵种模板
  void loadLogoTemplates();
  // 获取位置消息
  geometry_msgs::msg::PoseArray getArmorPosesMsg();
};